#!/usr/bin/env python

import numpy as np
import pylab as pl
import os.path as osp
import sys
import matplotlib
from matplotlib.font_manager import FontProperties as FP

# P-frame size estimator function
def p_est(r,fps):
	return r/float(fps)

# GOP size estimator function
def gop_est(r,fps,gop):
	return r/float(fps)*gop

# Search frame_log.txt file and build point list of P-frame sizes
def search_framelog(logfilep):

	bitrate_kB=[]
	p_av=[]
		
	line = '1'	
	while not line =='':
		line = logfilep.readline()
		if line.startswith('P-frame average size'):
			p_av.append(float(line[line.find('size')+4:-1]))
		if line.startswith('Bitrate'):
			bitrate_kB.append(float(line[line.find('Bitrate')+7:-5]))

	return bitrate_kB,p_av

# Search gop_log.txt file and build point list of gop sizes
def search_goplog(logfilep):

	gop_av=[]
		
	line = '1'	
	while not line =='':
		line = logfilep.readline()
		if line.startswith('Mean GOP size'):
			gop_av.append(float(line[line.find('size')+4:-3]))
	
	return gop_av

if __name__ == '__main__':

	# Use latex formatting in strings!
	matplotlib.rc('text', usetex=True) 

	# Font fix
	fp = FP()
	fp.set_size('small')

	# Parameters
	fps=24
	gop=60

	# Calculte p_est data
	p_est_x=np.arange(0.,10000.,10)
	p_est_y=[]	

	for bitrate in p_est_x:
		p_est_y.append(p_est(bitrate,fps))
	
	# Calculate gop_est data
	gop_est_x=np.arange(0.,10000.,10)
	gop_est_y=[]	

	for bitrate in gop_est_x:
		gop_est_y.append(gop_est(bitrate,fps,gop))

	# Prepare figure
	fig = pl.figure(figsize=[10,4])
	fig.subplots_adjust(wspace=0.5) # spacing between subplots
	ax1 = fig.add_subplot(121)
	ax2 = fig.add_subplot(122)

	# Fill estimator data into plot
	ax1.plot(p_est_x,p_est_y,label='$\hat{P}$')
	ax2.plot(gop_est_x,gop_est_y,label='$\hat{GOP}$')

	# Fill framelog data into plot
	framelog_filename = sys.argv[1]
	f = open(framelog_filename,'r')
	bitrate_p,size_p=search_framelog(f)
	f.close()

	ax1.plot(bitrate_p[0:3],size_p[0:3],marker='*',label='sintel\_1')
	ax1.plot(bitrate_p[3:6],size_p[3:6],marker='*',label='bbb\_1')
	ax1.plot(bitrate_p[6:9],size_p[6:9],marker='*',label='ed\_1')

	# Fill goplog data into plot
	goplog_filename=sys.argv[2]
	f = open(goplog_filename,'r')
	bitrate_gop=bitrate_p # Same videos!
	size_gop=search_goplog(f)

	ax2.plot(bitrate_gop[0:3],size_gop[0:3],marker='*',label='sintel\_1')
	ax2.plot(bitrate_gop[3:6],size_gop[3:6],marker='*',label='bbb\_1')
	ax2.plot(bitrate_gop[6:9],size_gop[6:9],marker='*',label='ed\_1')
		
	# Annotate plot
	ax1.set_ylabel('P-frame size [kB]')
	ax1.set_xlabel('Bit rate [kBps]')
	ax1.legend(prop=fp,loc='upper left')
	ax1.grid(True,which='both')
	ax2.set_ylabel('GOP size [kB]')
	ax2.set_xlabel('Bit rate [kBps]')
	ax2.legend(prop=fp,loc='upper left')
	ax2.grid(True,which='both')
	
	ax1.set_xlim(0,400)
	ax1.set_ylim(0,20)

	ax2.set_xlim(0,400)
	ax2.set_ylim(0,1100)

	# Save figure
	workingdir = osp.abspath('../../plotting/video_coding2')
	pl.savefig(workingdir+'/figs/estimator.eps')


